import sys

sys.path.append(".")

import argparse

from torch.utils.data import DataLoader
from auxmodels.core.runner import DevRunner
from auxmodels.core.utils import Configer
from auxmodels.core.worker import Trainer, Validator
from auxmodels.core.summarizer import Summarizer

import src.datasets.operations as op
from src.datasets.dataset import Augmenter
from src.datasets.dataset import CenterDataset
from src.models.model import FinetuneModel


def train(cfg_file):
    configer = Configer(cfg_file)
    cfg = configer.params
    print("cfg:", cfg)

    # 初始化模型
    model = FinetuneModel(cfg.hyper.learning_rate)

    # 数据增强相关设置
    augment = Augmenter(
        ops=[
            op.RandomRotateAndCrop(min_size=(512, 512), max_angle=30, select_mode="center"),
            op.RandomFlip(horizontal=0.5, vertical=0.5),
            op.Resize(input_size=(384, 384)),
        ]
    )
    # 训练数据集
    dataset_train = CenterDataset(cfg.dataset.train.file_list, augment, 1, (384, 384))
    dataloader_train = DataLoader(
        dataset_train,
        batch_size=cfg.dataset.train.batch_size,
        shuffle=True,
        num_workers=cfg.dataset.train.num_workers,
        pin_memory=True,
        collate_fn=dataset_train.collate_fn,
    )
    # 验证数据集
    dataset_val = CenterDataset(cfg.dataset.eval.file_list, augment, 1, (384, 384))
    dataloader_val = DataLoader(
        dataset_val,
        batch_size=cfg.dataset.eval.batch_size,
        shuffle=False,
        num_workers=cfg.dataset.eval.num_workers,
        pin_memory=True,
        collate_fn=dataset_val.collate_fn,
    )

    # 初始化训练器
    trainer = Trainer(dataloader_train)
    validator = Validator(dataloader_val, Summarizer())
    runner = DevRunner(
        model=model, configer=configer, gpu_id=0, trainer=trainer, validator=validator
    )
    # 启动训练
    runner.run()
    return


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfg", type=str)
    args = parser.parse_args()

    train(args.cfg)
